import sys
import os
filePath = sys.argv[1]
chain_list = []
aa_codes = {
    'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E',
    'PHE': 'F', 'GLY': 'G', 'HIS': 'H', 'LYS': 'K',
    'ILE': 'I', 'LEU': 'L', 'MET': 'M', 'ASN': 'N',
    'PRO': 'P', 'GLN': 'Q', 'ARG': 'R', 'SER': 'S',
    'THR': 'T', 'VAL': 'V', 'TYR': 'Y', 'TRP': 'W'
}
pdb_file = open("./{}".format(filePath), 'r')
pdb_lines = [i for i in pdb_file]
pdb_ca = []
pdb_AA = []
pdb_AA_all = []
seqs = []

for i in pdb_lines:
    if (i[0:5].strip() == "ATOM") & (i[13:16].strip() == "CA"):
        pdb_ca.append(i)
for i in pdb_ca:
    chain_list.append(i[21:23].strip())
chain_list = list(set(chain_list))

for i in chain_list:
    for j in pdb_ca:
        if j[21:23].strip() == i:
            pdb_AA.append(aa_codes[j[17:20].strip()])
    pdb_AA_all.append(pdb_AA)
    pdb_AA = []

for i in pdb_AA_all:
    seqs.append("".join(i))
print(chain_list)
print(seqs)

fasta_list = []
for i in range(len(chain_list)):
    fasta_list.append(">{}\n".format(chain_list[i]))
    fasta_list.append("{}\n".format(seqs[i]))

f = open("{}.fasta".format(filePath.split(".")[0]),"r")
f.writelines(fasta_list)
f.close()

